﻿#include "VulkanShaders.h"
#include <vulkan/vulkan.hpp>
#include "Render/VulkanContext.h"
#include "VulkanShader.h"
#include "Render/PreRender/PreRenderShader.h"
#include "Render/Material/WriteInShader.h"

std::shared_ptr<CVulkanShader> CVulkanShaders::CreateShaderFromRaw(const TVector<uint8_t> &InData, vk::ShaderStageFlagBits InStage)
{
    auto VulkanShader=std::make_shared<CVulkanShader>();
    vk::ShaderModuleCreateInfo ShaderInfo=vk::ShaderModuleCreateInfo()
        .setCodeSize(InData.Num())
        .setPCode((uint32_t*)InData.Data());

    VulkanShader->Module=CVulkanContext::Get()->GetDevice().createShaderModule(ShaderInfo);       
    VulkanShader->Stage=InStage;

    return VulkanShader;
}

CVulkanShaders::CVulkanShaders()
{
    CreateVulkanOnlyShaders();
}

void CVulkanShaders::UpdateShaders(const TVector<std::shared_ptr<CPreRenderShader>> &InShaders)
{
    auto OldShaders = Shaders;
    Shaders.Clear();

    TVector<TThreadSafeSharedPtr<CWriteInShader>> NeedUpdateShaderDatas;
    for (auto& PreRenderShader : InShaders)
    {
        auto Shader = PreRenderShader->WriteInData;
        auto FindShader = OldShaders.Find(Shader);
        if (FindShader)
        {
            Shaders.Add(Shader,*FindShader);
            if(Shader->IsDirty())
            {
                NeedUpdateShaderDatas.Add(Shader);
            }
            
        }
        else
        {
            Shaders.Add(Shader, std::make_shared<CVulkanShader>());
            NeedUpdateShaderDatas.Add(Shader);
        }
    }

    OldShaders.Clear();

    for (auto& ShaderData : NeedUpdateShaderDatas)
    {
        auto VulkanShader = Shaders[ShaderData];
        std::lock_guard<std::mutex> LockGuard(ShaderData->Mutex);
        ShaderData->bDirty = false;

        
        vk::ShaderModuleCreateInfo ShaderInfo=vk::ShaderModuleCreateInfo()
        .setCodeSize(ShaderData->Data.Num())
        .setPCode((uint32_t*)ShaderData->Data.Data());

        //回收旧的ShaderModule
        if(VulkanShader->Module)
        {
            CVulkanContext::Get()->GetRecycleObjects()->ShaderModules.Add(VulkanShader->Module);
            VulkanShader->Module=nullptr;
        }
        
        VulkanShader->Module=CVulkanContext::Get()->GetDevice().createShaderModule(ShaderInfo);       

        switch(ShaderData->Stage)
        {
        case EShaderStage::Vertex:
            VulkanShader->Stage=vk::ShaderStageFlagBits::eVertex;
            break;
        case EShaderStage::Fragment:

            VulkanShader->Stage=vk::ShaderStageFlagBits::eFragment;
            break;

        default:
            assert(false && "Unknow Shader Type");
            break;
        }

    }
}

std::shared_ptr<CVulkanShader> CVulkanShaders::FindVulkanObj(TThreadSafeSharedPtr<CWriteInShader> Shader)
{
    return Shaders[Shader];
}

std::shared_ptr<CVulkanShader> CVulkanShaders::GetVulkanOnlyShader(const CString &Name)
{
    return VulkanOnlyShaders[Name];
}

void CVulkanShaders::RegisterVulkanOnlyShader(const CString &Name, std::shared_ptr<CVulkanShader> Shader)
{
    VulkanOnlyShaders.Add(Name,Shader);
}

void CVulkanShaders::CreateVulkanOnlyShaders()
{
    //add RectWidget Vertex Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/RectWidget/RectWidget.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"RectWidgetVertex",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }
    
    //add RectWidget Fragment Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/RectWidget/RectWidget.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"RectWidgetFragment",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add ImageWidget Vertex Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/ImageWidget/ImageWidget.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"ImageWidgetVertex",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add ImageWidget Fragment Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/ImageWidget/ImageWidget.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"ImageWidgetFragment",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add CanvasWidgetLine Vertex Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/CanvasWidget/CanvasWidgetLine.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"CanvasWidgetLineVertex",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add CanvasWidgetLine Fragment Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/CanvasWidget/CanvasWidgetLine.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"CanvasWidgetLineFragment",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add CanvasWidgetTriangle Vertex Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/CanvasWidget/CanvasWidgetTriangle.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"CanvasWidgetTriangleVertex",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add CanvasWidgetTriangle Fragment Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/CanvasWidget/CanvasWidgetTriangle.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"CanvasWidgetTriangleFragment",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add DefferShading Vertex Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/DefferRenderingComposition.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"DefferShadingVertex",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add SpotLightDefferShading Fragment Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Lights/SpotLightDefferShading.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"SpotLightDefferShadingFragment",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add PointLightDefferShading Fragment Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Lights/PointLightDefferShading.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"PointLightDefferShadingFragment",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add DirectionalLightDefferShading Fragment Shader
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Lights/DirectionalLightDefferShading.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"DirectionalLightDefferShadingFragment",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }



    //add Border_LeftBottom.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/RectWidget/Border_LeftBottom.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"Border_LeftBottom_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add Border_LeftTop.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/RectWidget/Border_LeftTop.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"Border_LeftTop_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add Border_RightBottom.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/RectWidget/Border_RightBottom.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"Border_RightBottom_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add Border_RightTop.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/RectWidget/Border_RightTop.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"Border_RightTop_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add Rect_LeftBottom.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/RectWidget/Rect_LeftBottom.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"Rect_LeftBottom_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add Rect_LeftTop.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/RectWidget/Rect_LeftTop.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"Rect_LeftTop_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add Rect_RightBottom.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/RectWidget/Rect_RightBottom.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"Rect_RightBottom_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add Rect_RightTop.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/RectWidget/Rect_RightTop.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"Rect_RightTop_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add CircleWidget.vert
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Circle/CircleWidget.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"CircleWidget_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add CircleWidget.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Circle/CircleWidget.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"CircleWidget_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add RingWidget.vert
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Circle/RingWidget.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"RingWidget_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add RingWidget.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Circle/RingWidget.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"RingWidget_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add CanvasWidgetImage.vert
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/CanvasWidget/CanvasWidgetImage.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"CanvasWidgetImage_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add CanvasWidgetImage.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/CanvasWidget/CanvasWidgetImage.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"CanvasWidgetImage_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add CanvasWidgetTextImage.vert
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/CanvasWidget/CanvasWidgetTextImage.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"CanvasWidgetTextImage_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add CanvasWidgetTextImage.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/CanvasWidget/CanvasWidgetTextImage.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"CanvasWidgetTextImage_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add CanvasWidgetCircle.vert
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/CanvasWidget/CanvasWidgetCircle.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"CanvasWidgetCircle_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add CanvasWidgetCircle.frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/CanvasWidget/CanvasWidgetCircle.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"CanvasWidgetCircle_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add Axis vert
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Axis/Axis.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"Axis_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }
    //add Axis frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Axis/Axis.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"Axis_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add PointLightShadowMap vert
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Lights/PointLightShadowMap.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"PointLightShadowMap_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add PointLightShadowMap frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Lights/PointLightShadowMap.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"PointLightShadowMap_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add DirectionalLightShadowMap vert
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Lights/DirectionalLightShadowMap.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"DirectionalLightShadowMap_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add ambientlight vert
    {
        TVector<uint8_t> ShaderData={
        #include "Shaders/Lights/AmbientLightDefferShading.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"AmbientLight_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add ambientlight frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/Lights/AmbientLightDefferShading.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"AmbientLight_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add AntiAliasing vert
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/AntiAlising/FXAA.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"FXAA_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add AntiAliasing frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/AntiAlising/FXAA.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"FXAA_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add TextWidget vert
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/ImageWidget/TextWidget.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"TextWidget_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add TextWidget frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/ImageWidget/TextWidget.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"TextWidget_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

    //add StaticTextWidget vert
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/ImageWidget/StaticTextWidget.vert.spv.h"
        };
        VulkanOnlyShaders.Add(U"StaticTextWidget_vert",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eVertex));
    }

    //add StaticTextWidget frag
    {
        TVector<uint8_t> ShaderData={
            #include "Shaders/ImageWidget/StaticTextWidget.frag.spv.h"
        };
        VulkanOnlyShaders.Add(U"StaticTextWidget_frag",CreateShaderFromRaw(ShaderData,vk::ShaderStageFlagBits::eFragment));
    }

}

